/*
 * HSM Proxy Project.
 * Copyright (C) 2013 FedICT.
 * Copyright (C) 2013 Frank Cornelis.
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License version
 * 3.0 as published by the Free Software Foundation.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, see 
 * http://www.gnu.org/licenses/.
 */

package be.fedict.hsm.model;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.security.InvalidKeyException;
import java.security.KeyStore.PrivateKeyEntry;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.Signature;
import java.security.SignatureException;
import java.security.cert.Certificate;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import javax.annotation.PostConstruct;
import javax.ejb.ConcurrencyManagement;
import javax.ejb.ConcurrencyManagementType;
import javax.ejb.EJB;
import javax.ejb.Lock;
import javax.ejb.LockType;
import javax.ejb.Singleton;
import javax.ejb.Startup;
import javax.persistence.EntityManager;
import javax.persistence.PersistenceContext;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import be.fedict.hsm.entity.KeyStoreEntity;

@Singleton
@Startup
@ConcurrencyManagement(ConcurrencyManagementType.CONTAINER)
public class KeyStoreSingletonBean {

	private static final Log LOG = LogFactory
			.getLog(KeyStoreSingletonBean.class);

	private static final byte[] SHA1_DIGEST_INFO_PREFIX = new byte[] { 0x30,
			0x21, 0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05,
			0x00, 0x04, 0x14 };

	private static final byte[] SHA256_DIGEST_INFO_PREFIX = new byte[] { 0x30,
			0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, (byte) 0x86, 0x48, 0x01, 0x65,
			0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20 };

	private static final byte[] SHA512_DIGEST_INFO_PREFIX = new byte[] { 0x30,
			0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, (byte) 0x86, 0x48, 0x01, 0x65,
			0x03, 0x04, 0x02, 0x03, 0x05, 0x00, 0x04, 0x40 };

	private static final Map<String, byte[]> digestInfoPrefixes;

	static {
		digestInfoPrefixes = new HashMap<String, byte[]>();
		digestInfoPrefixes.put("SHA-1", SHA1_DIGEST_INFO_PREFIX);
		digestInfoPrefixes.put("SHA-256", SHA256_DIGEST_INFO_PREFIX);
		digestInfoPrefixes.put("SHA-512", SHA512_DIGEST_INFO_PREFIX);
	}

	/**
	 * key store id -> key store alias -> private key
	 */
	private Map<Long, Map<String, PrivateKeyEntry>> privateKeyEntries;

	@PersistenceContext
	private EntityManager entityManager;

	@EJB
	private KeyStoreLoader keyStoreLoader;

	@PostConstruct
	@Lock(LockType.WRITE)
	public void loadKeys() {
		LOG.debug("load keys...");
		this.privateKeyEntries = new HashMap<Long, Map<String, PrivateKeyEntry>>();
		List<KeyStoreEntity> keyStoreEntities = KeyStoreEntity
				.getList(this.entityManager);
		for (KeyStoreEntity keyStoreEntity : keyStoreEntities) {
			load(keyStoreEntity);
		}
	}

	private boolean load(KeyStoreEntity keyStoreEntity) {
		Map<String, PrivateKeyEntry> entries = this.keyStoreLoader
				.loadKeyStore(keyStoreEntity);
		if (null != entries) {
			this.privateKeyEntries.put(keyStoreEntity.getId(), entries);
			return true;
		} else {
			this.privateKeyEntries.remove(keyStoreEntity.getId());
			return false;
		}
	}

	/**
	 * Sign the given digest value.
	 * 
	 * @param keyStoreId
	 * @param keyStoreAlias
	 * @param digestAlgo
	 * @param digestValue
	 * @return the signature, or <code>null</code> in case something went wrong.
	 * @throws NoSuchAlgorithmException
	 * @throws InvalidKeyException
	 * @throws IOException
	 * @throws SignatureException
	 */
	@Lock(LockType.READ)
	public byte[] sign(long keyStoreId, String keyStoreAlias,
			String digestAlgo, byte[] digestValue)
			throws NoSuchAlgorithmException, InvalidKeyException, IOException,
			SignatureException {
		Map<String, PrivateKeyEntry> keyStoreKeys = this.privateKeyEntries
				.get(keyStoreId);
		if (null == keyStoreKeys) {
			LOG.error("unknown key store: " + keyStoreId);
			return null;
		}
		PrivateKeyEntry privateKeyEntry = keyStoreKeys.get(keyStoreAlias);
		if (null == privateKeyEntry) {
			LOG.error("private key for alias not available: " + keyStoreAlias);
			return null;
		}
		PrivateKey privateKey = privateKeyEntry.getPrivateKey();
		Signature signature = Signature.getInstance("NONEwithRSA");
		signature.initSign(privateKey);

		ByteArrayOutputStream digestInfo = new ByteArrayOutputStream();
		byte[] digestInfoPrefix = digestInfoPrefixes.get(digestAlgo);
		if (null == digestInfoPrefix) {
			throw new NoSuchAlgorithmException(digestAlgo);
		}
		digestInfo.write(digestInfoPrefix);
		digestInfo.write(digestValue);

		signature.update(digestInfo.toByteArray());

		return signature.sign();
	}

	@Lock(LockType.WRITE)
	public boolean newKeyStore(long keyStoreId) {
		KeyStoreEntity keyStoreEntity = this.entityManager.find(
				KeyStoreEntity.class, keyStoreId);
		LOG.debug("new key store: " + keyStoreId);
		return load(keyStoreEntity);
	}

	@Lock(LockType.WRITE)
	public List<String> getKeyStoreAliases(long keyStoreId) {
		Map<String, PrivateKeyEntry> keyStorePrivateKeys = this.privateKeyEntries
				.get(keyStoreId);
		if (null == keyStorePrivateKeys) {
			return new LinkedList<String>();
		}
		List<String> aliases = new LinkedList<String>();
		for (String alias : keyStorePrivateKeys.keySet()) {
			LOG.debug("key store alias: " + alias);
			aliases.add(alias);
		}
		return aliases;
	}

	@Lock(LockType.WRITE)
	public void removeKeyStore(long keyStoreId) {
		this.privateKeyEntries.remove(keyStoreId);
	}

	@Lock(LockType.WRITE)
	public boolean reload(long keyStoreId) {
		this.privateKeyEntries.remove(keyStoreId);
		KeyStoreEntity keyStoreEntity = this.entityManager.find(
				KeyStoreEntity.class, keyStoreId);
		return load(keyStoreEntity);
	}

	@Lock(LockType.WRITE)
	public Certificate[] getCertificateChain(long keyStoreId,
			String keyStoreAlias) {
		Map<String, PrivateKeyEntry> keyStorePrivateKeys = this.privateKeyEntries
				.get(keyStoreId);
		if (null == keyStorePrivateKeys) {
			LOG.error("key store not found: " + keyStoreId);
			return null;
		}
		PrivateKeyEntry privateKeyEntry = keyStorePrivateKeys
				.get(keyStoreAlias);
		if (null == privateKeyEntry) {
			LOG.error("no key entry found for alias: " + keyStoreAlias);
			return null;
		}
		return privateKeyEntry.getCertificateChain();
	}
}
